#!d:\anaconda3\envs\office_automation\python.exe -u
# -*- coding=utf-8 -*-
r"""
author: git config IVEN_CN && git config 13377529851@QQ.com
Date: 2024-01-17 11:15:15 +0800
LastEditTime: 2024-08-22 14:28:05 +0800
FilePath: \office_automation\color_div\init.py
details: 用于初始化颜色阈值，将颜色阈值保存成npy文件,识别的图片必须是800*1200的图片

Copyright (c) 2024 by IVEN, All Rights Reserved. 
"""
import sys
import os
import PIL.Image
import PIL.ImageFile
import cv2
import numpy as np
import PyQt5.QtWidgets as qw
from PyQt5.QtGui import QPixmap, QImage

PIL.ImageFile.LOAD_TRUNCATED_IMAGES = True  # 处理图像文件损坏的问题


def get_file_path(file_dir) -> list:
    """获取相对路径下文件夹的所有文件的路径

    :param file_dir: 文件夹路径"""
    file_path = []
    for root, dirs, files in os.walk(file_dir):
        for file in files:
            file_path.append(os.path.join(root, file))
    return file_path


class InitWindows(qw.QWidget):
    """用于初始化颜色阈值的窗口"""

    L_H = 0
    U_H = 180
    L_S = 0
    U_S = 255
    L_V = 0
    U_V = 255
    img = None
    cloth_area = 0
    model_area = 0
    _area = cloth_area
    index = 0
    COLOR_LST = [
        "白色2100",
        "黑色2010",
        "杏色2170",
        "粉红2340",
        "红色2200",
        "黄色2400",
        "深灰2040",
        "卡其2970",
        "克莱因蓝2500",
        "彩兰2670",
        "花灰2080",
        "灰蓝2690",
    ]
    COLOR_dst = {}
    color = COLOR_LST[0]
    mode = "cloth"

    def __init__(self, _folder_path):
        """
        初始化窗口

        :param _folder_path: 图片文件夹路径
        """
        self.folder_path = _folder_path
        self.COLOR_dst = {i: [[0, 0, 0], [180, 255, 255]] for i in self.COLOR_LST}
        super(InitWindows, self).__init__()
        self.img_paths = get_file_path(_folder_path)
        self.setWindowTitle("颜色阈值初始化")

        self.main_layout = qw.QHBoxLayout(self)  # 左侧阈值滑块，右侧显示图片
        self.setLayout(self.main_layout)

        # region 左侧
        self.left_widget = qw.QWidget()
        self.left_widget.setMinimumWidth(400)  # 设置最小宽度
        self.left_widget.setMaximumWidth(400)  # 设置最大宽度
        self.left_layout = qw.QVBoxLayout(self.left_widget)
        self.main_layout.addWidget(self.left_widget)

        # region 创建滑动条
        Trackbar_index, layout_index = self.createTrackbar(
            "图片索引", 0, len(self.img_paths) - 1, self.changeIndex
        )
        self.Trackbar_L_H, layout_L_H = self.createTrackbar(
            "L_H", 0, 180, self.changeLH
        )
        self.Trackbar_U_H, layout_U_H = self.createTrackbar(
            "U_H", 180, 180, self.changeUH
        )
        self.Trackbar_L_S, layout_L_S = self.createTrackbar(
            "L_S", 0, 255, self.changeLS
        )
        self.Trackbar_U_S, layout_U_S = self.createTrackbar(
            "U_S", 255, 255, self.changeUS
        )
        self.Trackbar_L_V, layout_L_V = self.createTrackbar(
            "L_V", 0, 255, self.changeLV
        )
        self.Trackbar_U_V, layout_U_V = self.createTrackbar(
            "U_V", 255, 255, self.changeUV
        )
        self.Trackbar_area, layout_area = self.createTrackbar(
            "面积", 0, 600 * 400, self.changeArea
        )
        self.left_layout.addLayout(layout_index)
        self.left_layout.addLayout(layout_L_H)
        self.left_layout.addLayout(layout_U_H)
        self.left_layout.addLayout(layout_L_S)
        self.left_layout.addLayout(layout_U_S)
        self.left_layout.addLayout(layout_L_V)
        self.left_layout.addLayout(layout_U_V)
        self.left_layout.addLayout(layout_area)
        # endregion

        # region 模式选择，平铺或模特
        mode_layout = qw.QVBoxLayout()
        mode_box = qw.QGroupBox("模式")
        cloth = qw.QRadioButton("平铺")
        model = qw.QRadioButton("模特")
        mode_box.setLayout(mode_layout)
        mode_layout.addWidget(cloth)
        mode_layout.addWidget(model)

        cloth.setChecked(True)

        cloth.toggled.connect(self.on_radio_button_toggled)
        model.toggled.connect(self.on_radio_button_toggled)
        self.left_layout.addWidget(mode_box)
        # endregion

        # region 下拉选项
        combo_box = qw.QComboBox()
        combo_box.addItems(self.COLOR_LST)
        combo_box.currentIndexChanged.connect(self.on_combobox_changed)
        self.left_layout.addWidget(combo_box)
        # endregion

        # region 保存按钮
        save_bt = qw.QPushButton("保存")
        save_bt.setFixedHeight(50)
        self.left_layout.addWidget(save_bt)
        save_bt.clicked.connect(self.save)
        # endregion

        # endregion

        # region 右侧
        self.right_layout = qw.QHBoxLayout()
        self.main_layout.addLayout(self.right_layout)

        # region 显示图片
        self.img = self.imread(self.img_paths[self.index])

        self.img_label = qw.QLabel()
        self.update_image_display()

        self.right_layout.addWidget(self.img_label)
        # endregion

        # region 显示掩膜
        self.filter()
        self.mask_Qmask = self.cvimg2Qimg(self.mask_img)

        self.mask_label = qw.QLabel()
        self.update_mask_display()

        self.right_layout.addWidget(self.mask_label)
        # endregion

        # endregion

        # 设置滑块的初始位置
        self.Trackbar_U_H.setValue(180)
        self.Trackbar_U_S.setValue(255)
        self.Trackbar_U_V.setValue(255)

        self.show()

    def update_threshold(self):
        """更新阈值"""
        self.COLOR_dst[self.color] = [
            [self.L_H, self.L_S, self.L_V],
            [self.U_H, self.U_S, self.U_V],
        ]

    def update_slider(self):
        """更新滑动条"""
        self.Trackbar_L_H.setValue(self.L_H)
        self.Trackbar_U_H.setValue(self.U_H)
        self.Trackbar_L_S.setValue(self.L_S)
        self.Trackbar_U_S.setValue(self.U_S)
        self.Trackbar_L_V.setValue(self.L_V)
        self.Trackbar_U_V.setValue(self.U_V)

    def createTrackbar(self, name: str, value, max_value, callback):
        """创建滑动条

        :param name: 滑动条名称
        :param value: 初始值
        :param max_value: 最大值
        :param callback: 回调函数
        """
        layout = qw.QHBoxLayout()
        label = qw.QLabel(name)
        layout.addWidget(label)
        slider = qw.QSlider()
        slider.setOrientation(1)
        slider.setValue(value)
        slider.setMaximum(max_value)
        slider.valueChanged.connect(callback)
        layout.addWidget(slider)
        return slider, layout

    def filter(self):
        if self.img is None:
            self.img = np.zeros((512, 512, 3), np.uint8)
        hsv = cv2.cvtColor(self.img, cv2.COLOR_BGR2HSV)
        low = np.array([self.L_H, self.L_S, self.L_V])
        up = np.array([self.U_H, self.U_S, self.U_V])
        self.mask_img = cv2.inRange(hsv, low, up)
        # contours是轮廓集，hierarchy是轮廓属性
        contours, hierarchy = cv2.findContours(
            self.mask_img,  # 二值图像
            cv2.RETR_TREE,  # 轮廓检索模式
            cv2.CHAIN_APPROX_SIMPLE,
        )  # 轮廓近似方法
        self.mask_img = cv2.cvtColor(self.mask_img, cv2.COLOR_GRAY2BGR)
        self.mask_img = cv2.resize(self.mask_img, (400, 600))
        for contour in contours:
            # 对每个轮廓进行矩形拟合
            x, y, w, h = cv2.boundingRect(contour)
            brcnt = np.array([[[x, y]], [[x + w, y]], [[x + w, y + h]], [[x, y + h]]])
            if w * h >= self.area:
                cv2.drawContours(
                    self.mask_img, [brcnt], -1, (0, 255, 0), 2
                )  # 在掩膜图像上绘制轮廓

    # region 回调函数
    def changeLH(self, value):
        self.L_H = value
        self.filter()
        self.update_mask_display()
        self.update_threshold()

    def changeUH(self, value):
        self.U_H = value
        self.filter()
        self.update_mask_display()
        self.update_threshold()

    def changeLS(self, value):
        self.L_S = value
        self.filter()
        self.update_mask_display()
        self.update_threshold()

    def changeUS(self, value):
        self.U_S = value
        self.filter()
        self.update_mask_display()
        self.update_threshold()

    def changeLV(self, value):
        self.L_V = value
        self.filter()
        self.update_mask_display()
        self.update_threshold()

    def changeUV(self, value):
        self.U_V = value
        self.filter()
        self.update_mask_display()
        self.update_threshold()

    def changeArea(self, value):
        self.area = value
        self.filter()
        self.update_mask_display()

    def changeIndex(self, value):
        self.index = value
        self.img = self.imread(self.img_paths[value])
        self.filter()
        self.update_image_display()
        self.update_mask_display()

    def update_image_display(self):
        self.Qimg = self.cvimg2Qimg(self.img)
        self.img_map = QPixmap.fromImage(self.Qimg)
        self.img_label.setPixmap(self.img_map)
        self.img_label.repaint()

    def update_mask_display(self):
        self.mask_Qmask = self.cvimg2Qimg(self.mask_img)
        self.mask_map = QPixmap.fromImage(self.mask_Qmask)
        self.mask_label.setPixmap(self.mask_map)
        self.mask_label.repaint()

    def on_radio_button_toggled(self):
        sender = self.sender()
        if isinstance(sender, qw.QRadioButton) and sender.isChecked():
            if sender.text() == "平铺":
                self.mode = "cloth"
                self.model_area = self.area
                self.area = self.cloth_area
                self.Trackbar_area.setValue(self.area)
            elif sender.text() == "模特":
                self.mode = "model"
                self.cloth_area = self.area
                self.area = self.model_area
                self.Trackbar_area.setValue(self.area)
            self.filter()
            self.update_mask_display()

    def on_combobox_changed(self, index):
        self.color = self.COLOR_LST[index]
        self.L_H, self.L_S, self.L_V = self.COLOR_dst[self.color][0]
        self.U_H, self.U_S, self.U_V = self.COLOR_dst[self.color][1]
        self.update_slider()

    def save(self):
        np.save(
            f"{self.folder_path}/{self.color}.npy",
            np.array([[self.L_H, self.L_S, self.L_V], [self.U_H, self.U_S, self.U_V]]),
        )
        np.save(f"{self.folder_path}/area_{self.mode}.npy", self.area * 4)

    # endregion

    def imread(self, path):
        try:
            img = PIL.Image.open(path)
            img = np.asarray(img)
            img = cv2.resize(img, (400, 600))
            return img
        except Exception as e:
            print(f"无法读取图像文件 {path}: {e}")
            return np.zeros((400, 600, 3), np.uint8)  # 返回一个空白图像

    def cvimg2Qimg(self, img):
        shape = img.shape
        bytes_per_line = 3 * shape[1]
        return QImage(
            img.data, shape[1], shape[0], bytes_per_line, QImage.Format_RGB888
        )

    @property
    def area(self):
        return self._area

    @area.setter
    def area(self, value):
        self._area = value


if __name__ == "__main__":
    app = qw.QApplication(sys.argv)
    windows = InitWindows(r"D:\work-ps\43girl7-29\A549")
    sys.exit(app.exec_())
